#include "npdm_patterns.h"
#include "npdm_operators.h"

using namespace Npdm;

//-----------------------------------------------------------------------------------------------------------------------------------------------------------
// Restrictions on operator strings for fermions:
// (1) If there is a contiguous piece of the operator string with the same site index, then all the creation ops in that piece should be to the left.
//     e.g. for 3PDM (1+,1+,1,,2+,3,3) is invalid.
// (2) If there are more than two creation or destruction ops with the same site index, expectation value is zero.
//

bool Npdm_operators::is_valid_opstring( std::vector< std::pair<CD,int> > & opstring )
{

  for (int i=1; i < opstring.size(); i++ ) {
    // If site same, but moves from destruction to creation then invalid string
    if ( opstring[i].second == opstring[i-1].second ) {
      if ( (opstring[i-1].first == DESTRUCTION) 
          && (opstring[i].first == CREATION) ) {
//std::cout << "invalid\n";
//print_opstring( opstring );
        return false;
      }
    }
  }

  // Check that no more than two c/d ops with same index
  int k = 1;
  for (int i=1; i < opstring.size(); i++) {
    if ( (opstring[i].second == opstring[i-1].second) 
       && (opstring[i].first == opstring[i-1].first) ) {
      k++;
    }
    else {
      k=1;
    }
    if ( k > 2 ) {
//std::cout << "invalid\n";
//print_opstring( opstring );
      return false;
    }
  }

  return true;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_6index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    for (int j=i; j <= site2; j++) {
      for (int k=j; k <= site2; k++) {
        for (int l=k; l <= site2; l++) {
          for (int m=l; m <= site2; m++) {
            for (int n=m; n <= site2; n++) {
              std::vector< std::pair<CD,int> > opstring;
              opstring.push_back( std::make_pair( cdvec.at(0), i ) );
              opstring.push_back( std::make_pair( cdvec.at(1), j ) );
              opstring.push_back( std::make_pair( cdvec.at(2), k ) );
              opstring.push_back( std::make_pair( cdvec.at(3), l ) );
              opstring.push_back( std::make_pair( cdvec.at(4), m ) );
              opstring.push_back( std::make_pair( cdvec.at(5), n ) );
              if ( is_valid_opstring( opstring ) ) output.insert( opstring );
            }
          }
        }
      }
    }
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_5index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    for (int j=i; j <= site2; j++) {
      for (int k=j; k <= site2; k++) {
        for (int l=k; l <= site2; l++) {
          for (int m=l; m <= site2; m++) {
            std::vector< std::pair<CD,int> > opstring;
            opstring.push_back( std::make_pair( cdvec.at(0), i ) );
            opstring.push_back( std::make_pair( cdvec.at(1), j ) );
            opstring.push_back( std::make_pair( cdvec.at(2), k ) );
            opstring.push_back( std::make_pair( cdvec.at(3), l ) );
            opstring.push_back( std::make_pair( cdvec.at(4), m ) );
            if ( is_valid_opstring( opstring ) ) output.insert( opstring );
          }
        }
      }
    }
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_4index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    for (int j=i; j <= site2; j++) {
      for (int k=j; k <= site2; k++) {
        for (int l=k; l <= site2; l++) {
          std::vector< std::pair<CD,int> > opstring;
          opstring.push_back( std::make_pair( cdvec.at(0), i ) );
          opstring.push_back( std::make_pair( cdvec.at(1), j ) );
          opstring.push_back( std::make_pair( cdvec.at(2), k ) );
          opstring.push_back( std::make_pair( cdvec.at(3), l ) );
          if ( is_valid_opstring( opstring ) ) output.insert( opstring );
        }
      }
    }
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_3index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    for (int j=i; j <= site2; j++) {
      for (int k=j; k <= site2; k++) {
        std::vector< std::pair<CD,int> > opstring;
        opstring.push_back( std::make_pair( cdvec.at(0), i ) );
        opstring.push_back( std::make_pair( cdvec.at(1), j ) );
        opstring.push_back( std::make_pair( cdvec.at(2), k ) );
        if ( is_valid_opstring( opstring ) ) output.insert( opstring );
      }
    }
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_2index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    for (int j=i; j <= site2; j++) {
      std::vector< std::pair<CD,int> > opstring;
      opstring.push_back( std::make_pair( cdvec.at(0), i ) );
      opstring.push_back( std::make_pair( cdvec.at(1), j ) );
      if ( is_valid_opstring( opstring ) ) output.insert( opstring );
    }
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::set< std::vector< std::pair<CD,int> > > Npdm_operators::build_1index_ops(char block, int site1, int site2, std::vector<CD> & cdvec)
{
  std::set< std::vector< std::pair<CD,int> > > output;

  for (int i=site1; i <= site2; i++) {
    std::vector< std::pair<CD,int> > opstring(1, std::make_pair( cdvec.at(0), i ) );
    output.insert( opstring );
  }

  return output;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_operators::build_operator_set(char block, int dotsite, std::vector<CD> & cdvec)
{

  if (cdvec.size() == 1) {
    if (block == 'l')
      lhs_ops[cdvec] = build_1index_ops( block, 1, dotsite-1, cdvec );
    if (block == 'd')
      dot_ops[cdvec] = build_1index_ops( block, dotsite, dotsite, cdvec );
    if (block == 'r')
      rhs_ops[cdvec] = build_1index_ops( block, dotsite+1, nsites_, cdvec );
  }
  else if (cdvec.size() == 2) {
    if (block == 'l')
      lhs_ops[cdvec] = build_2index_ops( block, 1, dotsite-1, cdvec );
    if (block == 'd')
      dot_ops[cdvec] = build_2index_ops( block, dotsite, dotsite, cdvec );
    if (block == 'r')
      rhs_ops[cdvec] = build_2index_ops( block, dotsite+1, nsites_, cdvec );
  }
  else if (cdvec.size() == 3) {
    if (block == 'l')
      lhs_ops[cdvec] = build_3index_ops( block, 1, dotsite-1, cdvec );
    if (block == 'd')
      dot_ops[cdvec] = build_3index_ops( block, dotsite, dotsite, cdvec );
    if (block == 'r')
      rhs_ops[cdvec] = build_3index_ops( block, dotsite+1, nsites_, cdvec );
  }
  else if (cdvec.size() == 4) {
    if (block == 'l')
      lhs_ops[cdvec] = build_4index_ops( block, 1, dotsite-1, cdvec );
    if (block == 'd')
      dot_ops[cdvec] = build_4index_ops( block, dotsite, dotsite, cdvec );
    if (block == 'r')
      rhs_ops[cdvec] = build_4index_ops( block, dotsite+1, nsites_, cdvec );
  }
  else if (cdvec.size() == 5) {
    if (block != 'd') assert(false);
    dot_ops[cdvec] = build_5index_ops( block, dotsite, dotsite, cdvec );
  }
  else if (cdvec.size() == 6) {
    if (block != 'd') assert(false);
    dot_ops[cdvec] = build_6index_ops( block, dotsite, dotsite, cdvec );
  }
  else
    assert( false );

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_operators::print_opstring( std::vector< std::pair<CD,int> > opstring )
{
  char cd;
  for (auto op = opstring.begin(); op != opstring.end(); op++) {
    cd = '.';
    if (op->first == CREATION) cd = '+';
    std::cout << "(" << op->second << cd << ")";
  }
  std::cout << std::endl;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------
